package com.aikg.kgrag_java.service.impl;

import com.aikg.kgrag_java.pojo.dto.neo4j.Neo4jRetrievalDTO;
import com.aikg.kgrag_java.pojo.entity.ConversationInfo;
import com.aikg.kgrag_java.pojo.entity.IndexDocument;
import com.aikg.kgrag_java.pojo.entity.Prompts;
import com.aikg.kgrag_java.pojo.entity.Question;
import com.aikg.kgrag_java.pojo.enums.RoleEnum;
import com.aikg.kgrag_java.pojo.request.LLMRequestBody;
import com.aikg.kgrag_java.pojo.vo.ChunkInfoVO;
import com.aikg.kgrag_java.service.IEmbeddingService;
import com.aikg.kgrag_java.service.ILLMService;
import com.aikg.kgrag_java.service.INeo4jService;
import com.aikg.kgrag_java.service.QuestionService;
import com.aikg.kgrag_java.service.adapter.ElasticSearchSimilarityAdapter;
import com.aikg.kgrag_java.service.adapter.Neo4jRetrievalAdapter;
import com.github.benmanes.caffeine.cache.Cache;
import jakarta.annotation.Resource;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;

/**
 * @ClassName: RAGPipelineService
 * @Description: RAG问答流程管道的服务类
 * @Author: 李宗雨
 * @Date: 2025/1/5 20:03
 * @Version: 1.0
 **/
@Service
public class RAGPipelineService {

    // 保存对话的缓存
    @Resource(name = "conversationCache")
    Cache<String, ConversationInfo> conversationCache;

    @Autowired
    QuestionService questionService;

    //    向量化服务
    @Autowired
    public IEmbeddingService embeddingService;

    //    图数据库检索
    @Autowired
    public INeo4jService neo4jService;

    //    向量数据库检索
    @Autowired
    public IndexServiceImpl indexService;

    //    生成层服务
    @Autowired
    public ILLMService llmService;

    /**
     * 流式RAG问答
     *
     * @param query: 用户的问题
     * @param k: 前k个重要节点或文本向量块
     * @param projectId: 领域 ID
     * @return 流式输出回答
     * @throws IOException:
     */
    public SseEmitter rag(String query, int k, int projectId, String conversationId) throws IOException {
        // 因为前端存操作了，所以后端不需要异步保存问题
//        saveQuestion(query,projectId);

        ConversationInfo conversation = new ConversationInfo();
        List<Double> queryEmbedding = embeddingService.embeddingQuery(query);
        Neo4jRetrievalDTO neo4jRetrievalDTO = neo4jService.neo4jRetrieval(query, projectId, k);
        // 文本向量库检索
        List<IndexDocument> result = indexService.searchByCosineSimilarity(
                queryEmbedding,
                3,
                "" + projectId,
                "vector"
        );
        String contexts = ElasticSearchSimilarityAdapter.build(result);
        LLMRequestBody llmRequestBody = new LLMRequestBody();
        // 创建prompt并拼接成message
        llmRequestBody.addMessages(RoleEnum.USER, Prompts.getRagPrompt(neo4jRetrievalDTO, contexts, query));

        // 封装给前端的数据
        conversation.setQuery(query);
        conversation.setConversationId(conversationId);
        conversation.setNeo4jRetrievalVO(Neo4jRetrievalAdapter.build(neo4jRetrievalDTO));
        // 将对话信息保留至缓存中
        List<ChunkInfoVO> chunks = new ArrayList<>();
        for (int i=0; i < result.size(); i++) {
            IndexDocument temp =  result.get(i);
            chunks.add(
                    new ChunkInfoVO(
                            temp.getPageContent(),
                            temp.getSource(),
                            temp.getKnowledgedb(),
                            Integer.parseInt(temp.getChunkNumber())
                    ));
        }
        conversation.setChunks(chunks);
        conversationCache.put(conversationId, conversation);
        return llmService.streamChat(llmRequestBody, conversationId);
    }


    /**
     * 异步保存问题（没用，别看）
     * @param question: 问题
     * @param projectId: 领域id
     */
    @Async
    public void saveQuestion(String question, int projectId){
        questionService.saveOrUpdateOne(new Question(UUID.randomUUID().toString().replace("-", ""), question, 1, "" + projectId,null));
    }

}
